Nvidia Jetson NX 系列 02.搭建 tensorflow-gpu 環境

Posted on Edited on In

這篇承接上一個 opencv 安裝的內容，這次要繼續把 tensorflow 的環境弄好。

CUDA 安裝完成後，一樣要記得配置環境變量，之後在驅動 gpu 的時候會依照這個環境變量去抓取 libcudart.so.10 檔案。否則會報一個下面這樣的錯誤：

Could not load dynamic library ‘libcudart.so.10.0’

環境變量設置如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
$ vim ~/.bashrc

# 加入以下三行
export LD_LIBRARY_PATH=/usr/local/cuda/lib64
export PATH=$PATH:/usr/local/cuda/bin
export CUDA_HOME=$CUDA_HOME:/usr/local/cuda
# 設置完成後存檔

$ source ~/.bashrc

# 環境變量設置完成後就可以用指令查看 CUDA 版本囉
$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Wed_Oct_23_21:14:42_PDT_2019
Cuda compilation tools, release 10.2, V10.2.89

不過在這上面的安裝最麻煩的應該是曾聽到同事說每個內容物件都要注意版本號，稍一不相容就會報錯。要注意以下幾種東西的版本號：

  • tensorflow + 是否需要 GPU (1.15.2 之後的版本好像沒這個問題，等等會稍微提到)

  • CUDA 版本 nvcc --version

  • cuDNN 版本

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    $ sudo cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
    #define CUDNN_MAJOR 8
    #define CUDNN_MINOR 0
    #define CUDNN_PATCHLEVEL 0
    --
    #define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

    #endif /* CUDNN_VERSION_H */

    #版本 8.0.0 (CUDNN_MAJOR + MINOR + PATCH)

  • jetpack 版本 sudo apt-cache show nvidia-jetpack

  • Python 版本


圖片來源：tensorflow 官網安裝指南

安裝依賴包，若要在 virtualenv 安裝的話，記得還是切換到虛擬環境進行安裝。

1
2
3
4
5
6
$ sudo apt-get update
$ sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev
# Install and upgrade pip3.
$ sudo apt-get install python3-pip
$ sudo pip install -U pip testresources setuptools
$ sudo pip install -U numpy==1.16.1 future==0.17.1 mock==3.0.5 h5py==2.9.0 keras_preprocessing==1.0.5 keras_applications==1.0.8 gast==0.2.2 enum34 futures protobuf

最後安裝 tensorflow。

1
$ pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v44 tensorflow==1.15.2+nv20.06

這裡 v44 代表我的 jetpack 是 4.4 版本，安裝 tensorflow 1.15.2，nv 是 nvidia tensorflow container 版本，詳情可以看下圖。

圖片來源：nvidia deep learning frameworks docs

對照看上圖，不過最終我是直接進入 https://developer.download.nvidia.com/compute/redist/jp/v44 這個網址，直接看目前提供的 wheel 有沒有和我需求的版本號對的上。 v42 會區分 tensorflow 或 tensorflow-gpu，不過 v44 好像就沒這區別了，印象中爬其他文的時候，好像有看到在某個版本後若沒有對到 gpu 會自動降到 cpu 這樣的字眼。因此這裡我就直接用上了。


Reference


上一篇是 nvidia jetson nx 上使用 virtualenv 搭建 opencv-python 環境。